from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
import json,time
from typing import AsyncIterable, Callable, Awaitable, AsyncGenerator

async def create_single_chunk_iter(content: bytes) -> AsyncGenerator[bytes, None]:
    yield content

async def modify_stream_content(
        original_iterator: AsyncIterable[bytes],
        content_type: str,
        charset: str = "utf-8"
) -> AsyncGenerator[bytes, None]:
    async for chunk in original_iterator:
        process_start = time.time()
        try:
            text_content = chunk.decode(charset)
            modified_text = f"[Stream Modified] {text_content}"
            modified_chunk = modified_text.encode(charset)
        except UnicodeDecodeError:
            modified_chunk = chunk

        print(f"[Middleware] Processed chunk: {modified_chunk[:512]}... | Time cost: {time.time() - process_start:.4f}s", flush=True)
        yield modified_chunk

async def modify_non_stream_content(
        original_content: bytes,
        content_type: str,
        charset: str = "utf-8"
) -> bytes:
    print("mmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmmm")
    if content_type.startswith("application/json"):
        try:
            print("uuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuuu")
            print(original_content)
            print("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
            json_data = json.loads(original_content.decode(charset))
            print("jjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjjj")
            print(json_data)
            if isinstance(json_data, dict):
                print("kkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkkk")
                json_data["modified"] = True
                return json.dumps(json_data).encode(charset)
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass
    return original_content


class MinferRunMiddleware(BaseHTTPMiddleware):
    def __init__(self, app):
        super().__init__(app)

    async def dispatch(self, request: Request, call_next):
        print("-----------------engine-server middleware running-------------------")
        try:
            response = await call_next(request)
        except Exception as e:
            print("hahahahahahahahahahahahahahahahahhahahahahahahahahahahahhahahahahahahahahahhah")
            print(e)
        if response.status_code < 200 or response.status_code >= 300:
            print("88888888888888888888888888888888888888888888888888888")
            return response
        if request.url.path not in ["/v1/completions", "/v1/chat/completions"]:
            print("999999999999999999999999999999999999999999999999999999")
            return response
        content_type = response.headers.get("content-type", "")
        print("--------------------------Content type:", content_type)
        if "application/json" in content_type:
            print("33333333333333333333333333333333333333333333333333333333333333")
            original_content = b"".join([chunk async for chunk in response.body_iterator])
            modified_content = await modify_non_stream_content(original_content, content_type)
            response.body_iterator = create_single_chunk_iter(modified_content)
            response.headers["content-length"] = str(len(modified_content))
        elif "text/event-stream" in content_type:
            print("4444444444444444444444444444444444444444444444444444444444")
            response.body_iterator = modify_stream_content(
                response.body_iterator,
                content_type,
                response.charset
            )
            response.headers["transfer-encoding"] = "chunked"
        print("5555555555555555555555555555555555555555555555555555555555555555555555555555555")
        return response